package com.gitee.wsl.struct.tree.ext

import com.gitee.wsl.struct.tree.NodeElement
import com.gitee.wsl.struct.tree.Tree
import com.gitee.wsl.struct.tree.TreeElement

fun <T,E: TreeElement<T, E>> TreeElement<T, E>.treeToMap(nodeMap: MutableMap<T, TreeElement<T, E>> = mutableMapOf()){
    nodeMap[data] = this
    if(this is NodeElement) {
        children?.forEach {
            it.treeToMap(nodeMap)
        }
    }
}

fun <T,E: TreeElement<T, E>> NodeElement<T, E>.treeToMap():MutableMap<T, TreeElement<T, E>>{
    val nodeMap: MutableMap<T, TreeElement<T, E>> = mutableMapOf()
    children?.forEach {
        it.treeToMap(nodeMap)
    }
    return nodeMap
}